from grid_env_ideal_obs_repeat_task import *
from grid_agent import *
from checkpoint_utils import *
from maze_factory import *
from replay_config import *
import argparse
import json
import sys
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
from sklearn.manifold import TSNE
import random
from sklearn.decomposition import PCA
from matplotlib.animation import FuncAnimation
from sklearn.cluster import KMeans
import threading
from matplotlib.ticker import FuncFormatter

# analysis of phase space

def progress_bar(current, total, barLength = 100):
    percent = float(current) * 100 / total
    arrow = '-' * int(percent/100 * barLength - 1) + '>'
    spaces = ' ' * (barLength - len(arrow))

    print('Progress: [%s%s] %d %%' % (arrow, spaces, percent), end='\r')
    sys.stdout.flush()

@jit
def get_action(y):
    return jnp.argmax(y)
get_action_vmap = jax.vmap(get_action)

# 定义一个自定义的函数，用于计算任意维度的叉乘
def cross_product(a, b):
    # 确保输入是一维数组
    a = np.asarray(a).ravel()
    b = np.asarray(b).ravel()
    # 计算输出的维度
    n = len(a)
    # 初始化输出数组
    c = np.zeros(n)
    # 遍历每个维度，计算对应的叉乘分量
    for i in range(n):
        # 计算循环的起始和结束位置
        start = (i + 1) % n
        end = (i - 1) % n
        # 计算叉乘分量
        c[i] = a[start] * b[end] - a[end] * b[start]
    # 返回输出数组
    return c

def ivf():

    """ parse arguments
    """
    rpl_config = ReplayConfig()

    parser = argparse.ArgumentParser()
    parser.add_argument("--model_pth", type=str, default=rpl_config.model_pth)
    parser.add_argument("--map_size", type=int, default=rpl_config.map_size)
    parser.add_argument("--task_pth", type=str, default=rpl_config.task_pth)
    parser.add_argument("--log_pth", type=str, default=rpl_config.log_pth)
    parser.add_argument("--nn_size", type=int, default=rpl_config.nn_size)
    parser.add_argument("--nn_type", type=str, default=rpl_config.nn_type)
    parser.add_argument("--show_kf", type=str, default=rpl_config.show_kf)
    parser.add_argument("--visualization", type=str, default=rpl_config.visualization)
    parser.add_argument("--video_output", type=str, default=rpl_config.video_output)
    parser.add_argument("--life_duration", type=int, default=rpl_config.life_duration)
    parser.add_argument("--start_i", type=int, default=rpl_config.start_i)
    parser.add_argument("--end_i", type=int, default=rpl_config.end_i)

    args = parser.parse_args()

    rpl_config.model_pth = args.model_pth
    rpl_config.map_size = args.map_size
    rpl_config.task_pth = args.task_pth
    rpl_config.log_pth = args.log_pth
    rpl_config.nn_size = args.nn_size
    rpl_config.nn_type = args.nn_type
    rpl_config.show_kf = args.show_kf
    rpl_config.visualization = args.visualization
    rpl_config.video_output = args.video_output
    rpl_config.life_duration = args.life_duration
    rpl_config.start_i = args.start_i
    rpl_config.end_i = args.end_i

    """ load model
    """
    params = load_weights(rpl_config.model_pth)
    # 定义一个函数，用于生成随机权重
    def init_weights_r(key, shape):
        return jax.random.normal(key, shape)
    # 生成一个随机的 PRNGKey
    key = jax.random.PRNGKey(np.random.randint(0, 1000))
    # key = jax.random.PRNGKey(4512)
    # 使用 tree_map 遍历 params 对象，并使用 init_weights 函数生成随机权重
    random_params = jax.tree_map(lambda x: init_weights_r(key, x.shape), params)
    # params = random_params

    # get elements of params
    tree_leaves = jax.tree_util.tree_leaves(params)
    for i in range(len(tree_leaves)):
        print("shape of leaf ", i, ": ", tree_leaves[i].shape)

    bias1 = np.array(tree_leaves[0])
    mat1 = np.array(tree_leaves[1])
    print("mat1.shape: ", mat1.shape)
    print("bias1: ", bias1)
    bias2 = np.array(tree_leaves[2])
    mat2 = np.array(tree_leaves[3])

    mat_obs = np.array(tree_leaves[1])[rpl_config.nn_size:rpl_config.nn_size+9,:]
    mat_intr = np.array(tree_leaves[1])[0:rpl_config.nn_size,:]
    print("mat_obs.shape: ", mat_obs.shape)

    """ create agent
    """
    if rpl_config.nn_type == "vanilla":
        model = RNN(hidden_dims = rpl_config.nn_size)
    elif rpl_config.nn_type == "gru":
        model = GRU(hidden_dims = rpl_config.nn_size)

    # check if param fits the agent
    if rpl_config.nn_type == "vanilla":
        assert params["params"]["Dense_0"]["kernel"].shape[0] == rpl_config.nn_size + 10

    n_samples = 10000
    k1 = npr.randint(0, 1000000)
    rnn_state = model.initial_state_rnd(n_samples, k1)
    rnn_state_old = rnn_state.copy()
    diff = jnp.abs(rnn_state - rnn_state_old)
    rnn_state_old = rnn_state.copy()
    diff_norm = jnp.linalg.norm(diff, axis=1)
    diff_norm_old = diff_norm.copy()
    norm_std = diff_norm.copy()

    rnn_state_trajectory = []
    intr_field = []

    for t in range(rpl_config.life_duration):

        progress_bar(t, rpl_config.life_duration)

        intr_vector = np.dot(rnn_state, mat_intr) + bias1
        intr_vector = np.tanh(intr_vector)
        intr_field.append(intr_vector - rnn_state)

        rnn_state_trajectory.append(np.array(rnn_state).copy())

        """ model forward 
        """
        rnn_state = intr_vector
        
    print(rnn_state.shape)
    print(norm_std.shape)
    rnn_state_np = np.array(rnn_state)

    # 将 rnn_state_trajectory 展开成 rnn_state_np 的形状
    rnn_state_trajectory_np = np.array(rnn_state_trajectory)
    rnn_state_trajectory_np = rnn_state_trajectory_np.reshape(-1, rnn_state_trajectory_np.shape[-1])
    print("shape of rnn_state_trajectory_np: ", rnn_state_trajectory_np.shape)

    # compute actions
    rnn_actions0 = np.dot(rnn_state_trajectory_np, mat2) + bias2
    print("shape of rnn_actions: ", rnn_actions0.shape)
    actions0 = get_action_vmap(jnp.array(rnn_actions0))
    actions0 = np.array(actions0)
    print("shape of actions: ", actions0.shape)

    # 将 intr_field 展开成 rnn_state_np 的形状
    intr_field_np = np.array(intr_field)
    intr_field_np = intr_field_np.reshape(-1, intr_field_np.shape[-1])
    print("shape of intr_field_np: ", intr_field_np.shape)

    # 对 rnn_state_np 进行 PCA
    intrinsic_pca = PCA()
    intrinsic_pca.fit(rnn_state_trajectory_np)

    # 打印 variance ratio
    print(intrinsic_pca.explained_variance_ratio_)

    rnn_state_trajectory_np_pca = intrinsic_pca.transform(rnn_state_trajectory_np)

    print("shape of rnn_state_trajectory_np_pca: ", rnn_state_trajectory_np_pca.shape)

    pc1 = intrinsic_pca.components_[0]
    pc2 = intrinsic_pca.components_[1]

    C = mat2[:,0]
    C1 = mat2[:,1]
    C2 = mat2[:,2]
    C3 = mat2[:,3]

    # Calculate matrix M
    M = np.column_stack((pc1, pc2))

    # Calculate pseudo-inverse
    pinvM = np.linalg.pinv(M)

    # Find coefficients 
    coeffs = pinvM.dot(C)

    # Reconstruct C 
    C_recon = coeffs[0]*pc1 + coeffs[1]*pc2

    coeffs2 = pinvM.dot(C_recon)
    # Reconstruct C2
    C_recon2 = coeffs2[0]*pc1 + coeffs2[1]*pc2

    # Compare
    def cmp_vectors(v1, v2):
        # 求 C_recon 和 C 的夹角
        d_ = np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2))
        angle_in_degree = np.arccos(np.clip(d_, -1.0, 1.0)) * 180 / np.pi
        return angle_in_degree
    
    # 求 C_recon 和 C 的夹角
    print("angle between C_recon and C: ", cmp_vectors(C_recon, C))
    # 求 pc1 和 pc2 的夹角
    print("angle between pc1 and pc2: ", cmp_vectors(pc1, pc2))
    # 求 pc1 和 C 的夹角
    print("angle between pc1 and C: ", cmp_vectors(pc1, C))
    # 求 pc2 和 C 的夹角
    print("angle between pc2 and C: ", cmp_vectors(pc2, C))
    # 求 C_recon2 和 C_recon 的夹角
    print("angle between C_recon2 and C_recon: ", cmp_vectors(C_recon2, C_recon))
    # 求 C_recon 和 pc1 的夹角
    print("angle between C_recon and pc1: ", cmp_vectors(C_recon, pc1))
    # 求 C_recon 和 pc2 的夹角
    print("angle between C_recon and pc2: ", cmp_vectors(C_recon, pc2))

    # 求 C 和 C1 的夹角
    print("angle between C and C1: ", cmp_vectors(C, C1))
    # 求 C 和 C2 的夹角
    print("angle between C and C2: ", cmp_vectors(C, C2))
    # 求 C 和 C3 的夹角
    print("angle between C and C3: ", cmp_vectors(C, C3))


    # # 计算 pc1 和 pc2 的外积
    # pc1_cross_pc2 = cross_product(pc1, pc2)
    # # 打印结果
    # # print(pc1_cross_pc2)

    # print("dot product of pc1 and pc1_cross_pc2: ", np.dot(pc1, pc1_cross_pc2))
    # print("dot product of pc2 and pc1_cross_pc2: ", np.dot(pc2, pc1_cross_pc2))
    # print("dot product of pc2 and pc1: ", np.dot(pc2, pc1))

if __name__ == "__main__":
    
    ivf()